import streamlit as st
import pandas as pd
import numpy as np
import mlflow
import os,json 
from sklearn.metrics import confusion_matrix, classification_report
from mlflow.tracking import MlflowClient 
import matplotlib.pyplot as plt
from utils.ml_clf import * 
from utils.ml_reg import * 
import plotly.graph_objects as go 
from sklearn.metrics import roc_curve, roc_auc_score, confusion_matrix, accuracy_score, precision_score, recall_score, f1_score
from utils.ml_utils  import get_search_params ,opt_model ,get_rg_models

plt.rcParams['font.sans-serif'] = ['SimHei']  
plt.rcParams['axes.unicode_minus'] = False  



# 初始化MLflow
mlflow.set_tracking_uri("http://localhost:8095")  
client = MlflowClient(tracking_uri="http://localhost:8095") 

# 初始化会话状态变量 - 只在首次运行时设置默认值


def ml_app():
    # 创建标签页
    tab1, tab2, tab3 , tab4= st.tabs([" 数据上传与配置", " 模型选择与训练", "模型管理", "模型预测"]) 

    with tab1:
        #st.header("上传训练数据") 
        st.markdown("""
                    <div style="color:brown; font-size:20px; margin-bottom:20px;">
                    <strong>上传训练数据</strong>
                    </div>
                    """, unsafe_allow_html=True) 
        
        # 数据上传
        
        train_file = st.file_uploader("上传训练数据 (CSV格式)", type=["csv"])
        val_file = st.file_uploader("上传验证数据 (可选，CSV格式)", type=["csv"]) 

        if train_file is not None:
            # 读取数据
            train_data = pd.read_csv(train_file) 
            st.markdown("""
                <div style="color:brown; font-size:20px; margin-bottom:20px;"> 
                <strong>训练数据预览</strong>
                </div>
                """, unsafe_allow_html=True)  
            st.dataframe(train_data.head()) 
            
            # 列选择 
            all_columns = train_data.columns.tolist()  
            ts_column = st.selectbox("选择时间戳列 (可选,选择了的话，后续验证集没有上传则会按照时间前后顺序划分)", [None] + all_columns) 
            target_column = st.selectbox("选择目标列", all_columns)
            feature_columns = st.multiselect("选择特征列", [col for col in all_columns if col != target_column])

            task_type = st.selectbox("模型类型", ["regression", "binary", "multiclass"], 
                                            format_func=lambda x: {"regression": "回归", "binary": "二分类", "multiclass": "多分类"}[x])
                    
            # 多分类需要指定类别数 
            num_classes = None   
            if task_type == "multiclass": 
                num_classes = st.number_input("类别数量", 2, 100, 2)   

            
            if ts_column:
                train_data[ts_column] = pd.to_datetime(train_data[ts_column])   
                train_data = train_data.sort_values(by=ts_column).reset_index(drop=True)  
            
            X_train = train_data[feature_columns]
            y_train = train_data[target_column] 

        # 验证集设置 
        if val_file is not None: 
            val_data = pd.read_csv(val_file) 
            st.markdown("""
                <div style="color:brown; font-size:20px; margin-bottom:20px;">
                <strong>验证数据预览</strong>
                </div>
                """, unsafe_allow_html=True) 
            st.dataframe(val_data.head()) 
            use_split = False 
            X_val = train_data[feature_columns]
            y_val = train_data[target_column]   

        else: 
            if train_file is not None: 
                use_split = True 
                test_size = st.slider("从训练集拆分验证集比例", 0.1, 0.3, 0.2)  
                shuffle = False if ts_column else  True  
                #print(shuffle,ts_column)
                #print(X_train)
                X_train, X_val, y_train, y_val = train_test_split(
                    X_train , y_train, test_size=0.2, random_state=42,shuffle=shuffle 
                    ) 
                #print(X_val)

        
        # 处理数据
        if st.button("处理数据"): 
            if not feature_columns:
                st.error("请至少选择一个特征列")
            else:
                with st.spinner("数据预处理中..."):
                    try:
                        
                        data_dict =  { 
                            'X_train': X_train, 'X_val': X_val,
                            'y_train': y_train, 'y_val': y_val,
                            'feature_names': feature_columns,
                            'target_column': target_column,
                            'num_classes':num_classes,
                            'task_type': task_type  # 
                        }  

                        if data_dict:
                            st.session_state.data_dict = data_dict 
                            st.success("数据预理完成！可以前往模型训练标签页进行训练")
                    except Exception as e:
                        st.error(f"数据预处理失败: {str(e)}")  
        else:
            if 'data_dict' not in st.session_state:
                st.session_state.data_dict = None  


    with  tab2:
        if st.session_state.data_dict is None:
            st.warning("⚠️ 请先在第一个标签页完成数据上传与配置")
        else:  
            # 选择模型类型 - 同样使用key直接绑定，不手动修改
            task_type = st.session_state.data_dict['task_type']
            
            col1, col2 = st.columns([2, 3.5]) 
            with col1:
                st.markdown("""
                <div style="color:brown; font-size:20px; margin-bottom:20px;">
                <strong>模型选择</strong>
                </div>
                """, unsafe_allow_html=True)  

                model_options = ["rf", "xgb", "lgb", "catboost"]
                model_display_names = {
                    "rf": "随机森林",
                    "xgb": "XGBoost",
                    "lgb": "LightGBM",
                    "catboost": "CatBoost"
                }   

                model_type =  st.selectbox(
                    "选择模型",
                    model_options,
                    format_func=lambda x: model_display_names[x],
                    key="model_type",  #  
                    index=0 
                ) 

                # 设置实验名称和运行名称
                exp_name = st.text_input("实验名称", st.session_state.model_type, key="experiment_name") 
                run_name = st.text_input("运行名称", f"{st.session_state.model_type}_run", key="run_name") 


                # 交叉验证折数
                cv = st.selectbox("交叉验证折数(选择交叉验证的话，验证集的时间顺序会失效)", [0,5,3,8], key="cv") 
                
                # 最大评估次数
                max_evals = st.slider("超参数最大搜索次数", 5, 100, 20, key="max_evals")   

            with col2: 
                # 根据模型类型和任务类型显示超参数配置
                st.markdown("""
                <div style="color:brown; font-size:20px; margin-bottom:20px;">
                <strong>超参数搜索范围配置</strong>
                </div>
                """, unsafe_allow_html=True)   
                search_space = get_search_params(model_type,task_type) 


            if st.button("开始训练",key='训练'):
                with st.spinner("正在训练..."):
                    #task_type = st.session_state.data_dict['task_type']
                    print('task_type',task_type,'**'*10,task_type)
                    try:
                        run_results = opt_model(st.session_state.data_dict ,task_type,model_type,search_space,max_evals,cv) 

                        # 显示训练完成信息
                        st.success("✅ 模型训练完成！") 
                        # 保存最佳参数
                        st.session_state.best_params = run_results 
                        col5,col6  = st.columns(2)
                        with col5:
                            st.write("最佳详细指标：") 
                            metrics =  run_results['metrics'].copy() 
                            for k in run_results['metrics'] .keys():
                                if 'val' not in k:
                                    del metrics[k] 
                            st.json(run_results['metrics']) 
                        with col6:
                            st.write("最佳参数详情：")  
                            del run_results['params']["features"]
                            del run_results['params']["target"]

                            st.json(run_results['params'])   

                    except Exception as e:
                        st.warning(f"训练出错: {str(e)}")    
            
    with tab3:
        external_website_url = "http://localhost:8095"  

        # 显示外部网站的内容
        st.markdown(f'<iframe src="{external_website_url}" width="100%" height="800"></iframe>', unsafe_allow_html=True)
            

    with  tab4:
        # 设置页面标题
        #st.title("模型预测")
        st.markdown("""
            <div style="color:brown; font-size:20px; margin-bottom:20px;">
            <strong>模型预测</strong>
            </div>
            """, unsafe_allow_html=True) 
        
        # 1. 上传测试数据集
        #st.subheader("1. 上传测试数据集")
        st.markdown("""
            <div style="color:black; font-size:18px; margin-bottom:20px;">
            <strong>上传测试数据集</strong>
            </div>
            """, unsafe_allow_html=True) 
        test_file = st.file_uploader("请上传CSV格式的测试数据", type=["csv"])
        
        if test_file is not None:
            # 读取测试数据
            test_data = pd.read_csv(test_file)
            st.success("测试数据上传成功！")
            st.dataframe(test_data.head())
            
            # 2. 展示注册模型列表
            #st.subheader("2. 选择预测模型")
            st.markdown("""
            <div style="color:black; font-size:18px; margin-bottom:20px;">
            <strong>选择预测模型</strong>
            </div>
            """, unsafe_allow_html=True) 
            model_df = get_rg_models()
            
            if not model_df.empty:
                st.dataframe(model_df) 

                model_options = [f"{row['模型名称']} ( 版本: {row['version_id']}, val_loss: {row['val_loss']} )" 
                            for _, row in model_df.iterrows()] 
        
                selected_model = st.selectbox("选择要使用的模型", model_options)

                # 3. 预测参数设置
                #st.subheader("3. 预测参数")
                st.markdown("""
                    <div style="color:black; font-size:18px; margin-bottom:20px;">
                    <strong>预测参数</strong> 
                    </div>
                    """, unsafe_allow_html=True)  
                
                # 获取数据字典信息（假设从模型元数据中获取，这里简化处理）
                # 实际应用中可能需要从模型注册中加载相关元数据
                if st.button("开始预测",key='测试'):
                    with st.spinner("正在加载模型并进行预测..."):
                        try:
                            # 解析选择的模型信息 
                            selected_model_name = selected_model.split(" ( 版本: ")[0]
                            selected_version = selected_model.split(" ( 版本: ")[1].split(",")[0]
                            run_id =  model_df[(model_df['version_id'] == selected_version)& (model_df['模型名称'] == selected_model_name)]['run_id'].values[0]

                            logged_model = f'runs:/{run_id}/model'
                            #loaded_model = mlflow.pyfunc.load_model(logged_model) 
                            
                            
                            try:
                                ## rf  logged_model 没有预测概率 
                                loaded_model = mlflow.sklearn.load_model(logged_model)  
                            except:
                                ## cat  等不能用sklearn 
                                try:
                                    loaded_model = mlflow.catboost.load_model(logged_model)  
                                except:
                                    try:
                                        loaded_model = mlflow.xgboost.load_model(logged_model)  
                                    except:
                                        loaded_model = mlflow.lightgbm.load_model(logged_model) 
                                
                            run = mlflow.get_run(run_id)  
                            run_params = run.data.params
                            features = json.loads(run_params['features'])
                            target =  run_params['target'] 
                            task_type =  run_params['task_type']

                            y_pred = loaded_model.predict(test_data[features])   
                            y_true = test_data[target].values 

                            result_df = pd.DataFrame({  
                                    '预测值': y_pred,
                                    '实际值':test_data[target]   
                                })
                            

                            y_proba = None 
                            if task_type  != 'regression':
                                y_proba = loaded_model.predict_proba(test_data[features]) 
                                if task_type=='binary':
                                    result_df['概率'] =  y_proba[:, 1]   
                                else :
                                    for i in range(y_proba.shape[1]):
                                        result_df[f'概率_{i}'] =  y_proba[:, i]

                            st.write ('预测详情')
                            c1,c2 =st.columns([2,3])
                            with c1:
                                st.dataframe(result_df)    

                            # 5. 计算并展示评估指标
                            st.markdown("""
                                <div style="color:black; font-size:18px; margin-bottom:20px;">
                                <strong>预测评估指标</strong>
                                </div>
                                """, unsafe_allow_html=True) 
                            if task_type == 'regression':
                                mse_score, mae_score, rmse_score, r2 = calculate_metrics_reg(y_true, y_pred) 
                            
                                metrics ={
                                    'mse_score':mse_score,
                                    'mae_score':mae_score,
                                    'rmse_score':rmse_score,
                                    'r2':r2 
                                }
                                print(metrics)

                                col9,col10 = st.columns([1,4]) 

                                with col9:
                                    st.write('预测指标')   
                                    st.json(metrics)  
                                
                                with  col10 :
                                    st.write('预测散点图') 
                                    fig, ax = plt.subplots(figsize=(10,3)) 
                                    ax.scatter(y_true, y_pred,edgecolors="white",linewidth=0.2,c='blue')
                                    ax.scatter(y_true, y_true,linewidth=0.3,c='red')
                                    ax.spines["right"].set_visible(False)
                                    ax.spines["top"].set_visible(False)
                                    ax.set_xlabel('实际值', fontsize=10)
                                    ax.set_ylabel('预测值', fontsize=10)
                                    #ax.grid(True, color="gray")  
                                    ax.legend(loc="upper left", frameon=False, fontsize=6)  
                                    ax.tick_params(axis="both", labelsize=6)  

                                    st.pyplot(fig) 
                            
                            if task_type !='regression':
                                accuracy, precision, recall, f1, auc =  calculate_metrics_clf(y_true, y_pred, y_proba ,task_type)

                                metrics ={
                                    'accuracy':accuracy,
                                    'precision':precision,
                                    'recall':recall,
                                    'f1':f1,
                                    'auc':auc, 
                                }  
                            
                            
                                if task_type=='multiclass':
                                    del metrics['auc'] 

                                col7,col8 = st.columns([1,3]) 

                                with col7:
                                    st.write('预测指标')   
                                    st.json(metrics)  
                                with  col8:
                                    if  task_type == 'binary':
                                        
                                        c1,c2 = st.columns(2)
                                        
                                        with c1:
                                            st.write('auc 曲线')
                                            
                                            fpr, tpr, thresholds = roc_curve(y_true, y_proba[:,1])
                                            auc_score = roc_auc_score(y_true, y_proba[:,1])

                                            # 寻找最佳阈值
                                            optimal_idx = np.argmax(tpr - fpr)
                                            optimal_threshold = thresholds[optimal_idx]

                                            # 绘制 ROC 曲线
                                            fig = go.Figure()
                                            fig.add_trace(go.Scatter(x=fpr, y=tpr, mode='lines', name='ROC Curve',
                                                                    line=dict(color='orange', width=2)))
                                            fig.add_trace(go.Scatter(x=[0, 1], y=[0, 1], mode='lines', name='随机分类',
                                                                    line=dict(dash='dash', color='gray')))
                                            # 标记最佳阈值点
                                            fig.add_trace(go.Scatter(x=[fpr[optimal_idx]], y=[tpr[optimal_idx]], mode='markers',
                                                                    name=f'最佳阈值 ({optimal_threshold:.4f})',
                                                                    marker=dict(size=12, color='red', symbol='star')))
                                            fig.update_layout(
                                                title='ROC Curve',
                                                xaxis_title='假正率 (FPR)',
                                                yaxis_title='真正率 (TPR)',
                                                template="plotly_white",
                                                hovermode="x unified"
                                            )
                                            st.plotly_chart(fig)
                                        with  c2:

                                            # 6. 使用最佳阈值计算分类指标
                                            st.write("基于最佳阈值的分类指标") 
                                            # 根据最佳阈值进行预测
                                            y_pred_optimal = (y_proba[:,1] >= optimal_threshold).astype(int)

                                            # 计算指标
                                            acc = accuracy_score(y_true, y_pred_optimal)
                                            prec = precision_score(y_true, y_pred_optimal)
                                            rec = recall_score(y_true, y_pred_optimal)
                                            f1 = f1_score(y_true, y_pred_optimal)
                                            #cm = confusion_matrix(y_true, y_pred_optimal)

                                            metrics_best ={
                                                'accuracy':acc,
                                                'precision':prec,
                                                'recall':rec,
                                                'f1':f1,
                                            }
                                            
                                            st.json(metrics_best) 

                                    else:
                                        fig = go.Figure()
                                        fig.add_shape(type='line', line=dict(dash='dash'), x0=0, x1=1, y0=0, y1=1)

                                        per_class_auc = []
                                        per_class_metrics = {
                                            '类别': list(range(y_proba.shape[1])),

                                            '最佳阈值': [],
                                            '精确率 (Precision)': [],
                                            '召回率 (Recall)': [],
                                            'F1 分数 (F1-score)': [],
                                            'AUC': []
                                        }
                                        for i in range(y_proba.shape[1]):
                                            y_true_binary = (y_true == i).astype(int)
                                            fpr, tpr, thresholds  = roc_curve(y_true_binary, y_proba[:, i])
                                            auc_score = roc_auc_score(y_true_binary, y_proba[:, i])
                                            
                                            
                                            # 找到最佳阈值 (Youden's J 统计量)
                                            optimal_idx = np.argmax(tpr - fpr)
                                            optimal_threshold = thresholds[optimal_idx]
                                            
                                            # 基于最佳阈值进行预测并计算指标
                                            y_pred_optimal = (y_proba[:, i] >= optimal_threshold).astype(int)
                                            prec = precision_score(y_true_binary, y_pred_optimal, zero_division=0)
                                            rec = recall_score(y_true_binary, y_pred_optimal, zero_division=0)
                                            f1 = f1_score(y_true_binary, y_pred_optimal, zero_division=0)
                                            
                                            # 存储指标
                                            per_class_metrics['最佳阈值'].append(f"{optimal_threshold:.4f}")
                                            per_class_metrics['精确率 (Precision)'].append(f"{prec:.4f}")
                                            per_class_metrics['召回率 (Recall)'].append(f"{rec:.4f}")
                                            per_class_metrics['F1 分数 (F1-score)'].append(f"{f1:.4f}")
                                            per_class_metrics['AUC'].append(f"{auc_score:.4f}")
        
                                            per_class_auc.append(auc_score)
                                            fig.add_trace(go.Scatter(x=fpr, y=tpr, mode='lines', name=f'ROC {i} (AUC = {auc_score:.4f})'))
                                            
                                        col1, col2 = st.columns([1, 1])

                                        with col1:
                                            st.write("ROC 曲线图 (一对多)")
                                            fig.update_layout(
                                                title='每个类别的 ROC 曲线',
                                                xaxis_title='假正率 (FPR)',
                                                yaxis_title='真正率 (TPR)',
                                                template="plotly_white",
                                                hovermode="x unified"
                                            )
                                            st.plotly_chart(fig, use_container_width=True)

                                        with col2:
                                            st.write("以下表格展示了每个类别的精确率、召回率、F1 分数和 AUC (一对多)。")
                                            
                                            
                                            # 7. 将所有指标放入一个 DataFrame
                                            metrics_df = pd.DataFrame(per_class_metrics) 
                                            
                                            st.dataframe(metrics_df, hide_index=True)

                                                                    

                            
                        
                        except Exception as e:
                                            
                            st.error(f"预测过程出错: {str(e)}") 


